# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: db\__init__.py
__author__ = "sanyi"
from settings_provider import settings
from sqlalchemy.pool import Pool
from sqlalchemy import exc as sqlalchemy_exc
from sqlalchemy import event as sqlalchemy_event
from sqlalchemy import create_engine as sqlalchemy_create_engine


def load_sqlalchemy_tables(create_tables=False, db_engine=None):
    from db.tables import metadata
    from db.tables.events import EventRow
    from db.tables.groups import GroupRow
    from db.tables.licenses import LicenseRow
    from db.tables.locations import LocationRow
    from db.tables.profiles import ProfileRow
    from db.tables.scan_sessions import ScanSessionRow
    from db.tables.scan_session_vulns import ScanSessionVulnRow
    from db.tables.scans import ScanRow
    from db.tables.target2group import Target2GroupRow
    from db.tables.target_vulns import TargetVulnRow
    from db.tables.targets import TargetRow
    from db.tables.user2shard import User2ShardRow
    from db.tables.user2group import User2GroupRow
    from db.tables.users import UserRow
    from db.tables.vuln_types import VulnTypeRow
    from db.tables.intents import IntentRow
    from db.tables.uploads import UploadRow
    from db.tables.scan_session_jobs import ScanSessionJobRow
    from db.tables.workers import WorkerRow
    from db.tables.reports import ReportRow
    from db.tables.report_templates import ReportTemplateRow
    from db.tables.scan_session_vulns_stats import ScanSessionVulnsStatsRow
    from db.tables.target_vulns_stats import TargetVulnsStatsRow
    from db.tables.targets_configurations import TargetConfigurationRow
    from db.tables.system_config import SystemConfigRow
    from db.tables.targets_allowed import TargetsAllowedRow
    from db.tables.issue_trackers import IssueTrackerRow
    from db.tables.phone_validation import PhoneValidationRow
    from db.tables.system_config_mu import SystemConfigMuRow
    from db.tables.pwd_history import PwdHistoryRow
    from db.tables.ui_sessions import UiSessionRow

    if create_tables:
        metadata.create_all(db_engine)


def populate_database(db_engine):
    load_sqlalchemy_tables(True, db_engine)


def ping_connection(dbapi_connection, _, connection_proxy):
    cursor = dbapi_connection.cursor()
    try:
        cursor.execute("SELECT 1")
    except Exception:
        connection_proxy._pool.dispose()
        raise sqlalchemy_exc.DisconnectionError()
    finally:
        try:
            cursor.close()
        except Exception:
            pass


def register_pessimistic_connection_handling():
    sqlalchemy_event.listen(Pool, "checkout", ping_connection)


__engine_pool = dict()
register_pessimistic_connection_handling()
import json
from helpers.json_extra import JsonEncoder


def json_serializer(data):
    return json.dumps(data, cls=JsonEncoder)


def get_engine_pooled(name="master", debug_mode=None):
    engine_key = name
    try:
        return __engine_pool[engine_key]
    except KeyError:
        descriptor = settings[("databases", "connections", name)]
        if debug_mode is None:
            debug_mode = settings.get("databases.debug", True)
        extra_arguments = descriptor.get("extra_arguments", dict())
        connection_string = (
            "postgresql+psycopg2://%(user)s:%(password)s@%(host)s:%(port)s/%(db)s"
            % descriptor["connection"]
        )
        db_engine = sqlalchemy_create_engine(
            echo=debug_mode,
            isolation_level="AUTOCOMMIT",
            json_serializer=json_serializer,
            **extra_arguments
        )
        __engine_pool[engine_key] = db_engine

    return db_engine


class Connection:
    def __init__(
        self,
        connection_name,
        create_transactions=False,
        transaction_isolation_level="READ COMMITTED",
        *args,
        **kwargs
    ):
        self._Connection__args = args
        self._Connection__kwargs = kwargs
        self._Connection__connection = None
        self._Connection__transaction_isolation_level = transaction_isolation_level
        self._Connection__connection_name = connection_name
        self._Connection__create_transactions = create_transactions
        self._Connection__transaction = None

    @property
    def name(self):
        return self._Connection__connection_name

    def __enter__(self):
        self._Connection__connection = (
            get_engine_pooled(
                self._Connection__connection_name,
                *self._Connection__args,
                **self._Connection__kwargs
            )
        ).connect()
        if not self._Connection__create_transactions:
            self._Connection__connection.execution_options(isolation_level="AUTOCOMMIT")
        else:
            self._Connection__connection.execution_options(
                isolation_level=self._Connection__transaction_isolation_level
            )
            self._Connection__transaction = self._Connection__connection.begin()
        return self._Connection__connection

    def __exit__(self, _exc_type, exc_val, exc_tb):
        if self._Connection__connection:
            pass
        if self._Connection__transaction:
            try:
                if _exc_type is None:
                    self._Connection__transaction.commit()
            except Exception:
                pass

            self._Connection__connection.close()
            self._Connection__connection = None


from psycopg2.extensions import adapt as sql_escape
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ColumnClause


class FunctionWrapper(ColumnClause):
    def params(self, *optional_dict, **kwargs):
        pass

    def unique_params(self, *optional_dict, **kwargs):
        pass

    def __init__(self, *args, **kwargs):
        self.kwargs = kwargs
        super().__init__(*args)


@compiles(FunctionWrapper)
def __compile_function_wrapper(element, compiler, **kw):
    s = (",").join(("%s:=%s" % (k, sql_escape(v)) for k, v in element.kwargs.items()))
    return "%s(%s)" % (element.name, s)
